import asyncio
from contextlib import asynccontextmanager

from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.prebuilt import create_react_agent

from mcp_agent.my_llm import llm

# mcp_server_config = {
#     "url": "http://localhost:8000/sse",
#     "transport": "sse"
# }

mcp_server_config = {  # 连接MCP服务器的配置
    "url": "http://127.0.0.1:8000/streamable",
    "transport": "streamable_http"  # MCP的通信机制
}


@asynccontextmanager
async def make_agent():
    """生成一个智能体(langgraph)"""
    # 创建客户端实例（不再使用async with）
    client = MultiServerMCPClient({'lx_mcp': mcp_server_config})

    try:
        # 使用推荐的get_tools方法获取工具（异步调用）
        tools = await client.get_tools()
        print("获取到的工具:", tools)

        # 创建智能体
        agent = create_react_agent(llm, tools=tools)
        yield agent
    finally:
        # 确保关闭客户端连接
        if hasattr(client, 'aclose'):
            await client.aclose()


async def main():
    """在异步环境下，创建智能体，并执行"""
    async with make_agent() as agent:
        # resp = await agent.ainvoke({'messages': '计算一下(3+6)的结果'})
        resp = await agent.ainvoke({'messages': '计算一下(3 + 5) x 12 - 2 的结果,计算的时候尽量调用工具,因为工具有特殊的业务逻辑。只返回数字结果。只返回工具计算的结果，不然将会收到惩罚'})
        # resp = await agent.ainvoke({'messages': '今天的北京的天气怎么样？'})
        print(resp.get('messages')[-1].content)


if __name__ == '__main__':
    asyncio.run(main())